import os.path
import shutil

import torch as t
from torch.utils.tensorboard import SummaryWriter
from tensorboard.backend.event_processing import event_accumulator
import matplotlib.pyplot as plt
import zipfile

def _init_logs_(config, train_loss, train_accuracy, test_loss, test_accuracy, epoch):
    log_dir = './log1/' + config.exp_name
    if epoch == 1 and os.path.exists(log_dir):
        shutil.rmtree(log_dir)
    writer = SummaryWriter(log_dir=log_dir)
    writer.add_scalar('Train Loss', train_loss, epoch)
    writer.add_scalar('Train Accuracy', train_accuracy, epoch)
    writer.add_scalar('Test Loss', test_loss, epoch)
    writer.add_scalar('Test Accuracy', test_accuracy, epoch)
    writer.close()


def tensorboard_read(config):
    path = './log1/' + config.exp_name
    ea = event_accumulator.EventAccumulator(path)
    ea.Reload()
    if os.path.exists('./tensorboard_fig/' + config.exp_name):
        shutil.rmtree('./tensorboard_fig/' + config.exp_name)
    os.makedirs('./tensorboard_fig/' + config.exp_name)
    for val_name in ea.scalars.Keys():
        val_acc = ea.scalars.Items(val_name)
        fig = plt.figure()
        plt.plot([i.step for i in val_acc], [j.value for j in val_acc], label=val_name)
        plt.title(val_name)
        plt.xlabel('step')
        plt.ylabel(val_name)
        plt.show()
        val_name = val_name.replace('/', '_')
        fig.savefig('./tensorboard_fig/' + config.exp_name + '/' + val_name + '.png')


def zip_my_log(config):
    path = './log1/' + config.exp_name
    with zipfile.ZipFile('log.zip', 'w', compression=zipfile.ZIP_LZMA) as f:
        f.write(path)
        f.close()

def print_model_param(model):
    try:
        os.makedirs('./param_fig')
    except:
        pass
    i = 0
    mean = []
    var = []
    for name, param in model.state_dict().items():
        if 'weight' in name and param.dim() == 4:
            mean.append(t.mean(abs(param)).item())
            var.append(t.var(param).item())
            i += 1
    fig = plt.figure()
    plt.plot(range(i), mean)
    plt.title('param_weight_mean')
    fig.savefig('./param_fig/param_weight_mean.png')
    fig = plt.figure()
    plt.plot(range(i), var)
    plt.title('param_weight_var')
    fig.savefig('./param_fig/param_weight_var.png')





